import cv2
import random
import numpy as np
from .util import _get_boarder, bbox_iou, bbox_crop, get_affine_transform, affine_transform
from .augment import RandomJitterAug


def HorizontalFlipAug(src, bbox, landms=None):
    width = src.shape[1]
    src = src[:, ::-1]
    bbox[:, ::2] = width - bbox[:, 2::-2]
    if landms is not None:
        landms[:, ::2] = width - landms[:, ::2]
        landms = landms[:, [2, 3, 0, 1, 4, 5, 8, 9, 6, 7]]
    return src, bbox, landms



def Safe_Crop(bbox, size):
    """Crop a random part of the input and rescale it to some size without loss of bbox"""
    w, h = size
    x1, y1 = min(bbox[:, 0]), min(bbox[:, 1])
    x2, y2 = max(bbox[:, 2]), max(bbox[:, 3])

    # find a biger region
    bx1, by1 = x1 * random.random(), y1 * random.random()
    bx2, by2 = x2 + (w - x2) * random.random(), y2 + (h - y2) * random.random()

    return [bx1, by1, bx2, by2]
    
def Random_Crop(size):
    """bbox is None"""
    w, h = size
    scale = random.uniform(0.3, 1.0)
    short_side = min(w, h)
    cw= int(scale*short_side)
    ch = cw
    x1, y1 = random.randrange(w-cw), random.randrange(h-ch)
    return [x1, y1, x1+cw, y1+ch]


def With_Constraints_Crop(bbox, size, min_scale=0.3, max_scale=1, 
        max_aspect_ratio=2, constraints=None, max_trial=50):
    """Crop an image randomly with bounding box constraints."""
    if constraints is None:
        constraints = (
            (0.1, None),
            (0.3, None),
            (0.5, None),
            (0.7, None),
            (0.9, None),
            (None, 1),
        )
    w, h = size
    candidates = [(0, 0, w, h)]
    for min_iou, max_iou in constraints:
        min_iou = -np.inf if min_iou is None else min_iou
        max_iou = np.inf if max_iou is None else max_iou

        for _ in range(max_trial):
            scale = random.uniform(min_scale, max_scale)
            aspect_ratio = random.uniform(
                max(1 / max_aspect_ratio, scale * scale),
                min(max_aspect_ratio, 1 / (scale * scale)))
            crop_h = int(h * scale / np.sqrt(aspect_ratio))
            crop_w = int(w * scale * np.sqrt(aspect_ratio))

            crop_t = random.randrange(h - crop_h)
            crop_l = random.randrange(w - crop_w)
            crop_bb = np.array((crop_l, crop_t, crop_l + crop_w, crop_t + crop_h))

            if len(bbox) == 0:
                top, bottom = crop_t, crop_t + crop_h
                left, right = crop_l, crop_l + crop_w
                return bbox, (left, top, right-left, bottom-top)

            iou = bbox_iou(bbox, crop_bb[np.newaxis])
            if min_iou <= iou.min() and iou.max() <= max_iou:
                top, bottom = crop_t, crop_t + crop_h
                left, right = crop_l, crop_l + crop_w
                candidates.append((left, top, right-left, bottom-top))
                break

    # random select one
    while candidates:
        crop = candidates.pop(np.random.randint(0, len(candidates)))
        new_bbox = bbox_crop(bbox, crop, allow_outside_center=False)
        if new_bbox.size < 1:
            continue
        new_crop = (crop[0], crop[1], crop[0]+crop[2], crop[1]+crop[3])
        return new_bbox, new_crop
    return bbox, (0, 0, w, h)




def CenterNet_Train_Crop(image, bbox, labels, out_size, landms=None, max_trial=50):
    h, w = image.shape[:2]
    s = max(h, w) * 1.0
    c = np.array([w / 2., h / 2.], dtype=np.float32)

    for _ in range(max_trial):
        c_bbox, c_labels = bbox.copy(), labels.copy()
        h_boarder = _get_boarder(out_size[1]//4, h)
        w_boarder = _get_boarder(out_size[0]//4, w)
        c[0] = np.random.randint(low=w_boarder, high=w - w_boarder)
        c[1] = np.random.randint(low=h_boarder, high=h - h_boarder)
        s = s * np.random.choice(np.arange(0.6, 1.4, 0.1))
        trans_input = get_affine_transform(c, s, 0, out_size)
        for i in range(c_bbox.shape[0]):
            c_bbox[i, :2] = affine_transform(c_bbox[i, :2], trans_input)
            c_bbox[i, 2:4] = affine_transform(c_bbox[i, 2:4], trans_input)
        centers = (c_bbox[:, :2] + c_bbox[:, 2:]) / 2
        # center in image
        mask = np.logical_and(0 < centers, centers < np.array([out_size[0], out_size[1]])).all(axis=1)

        b_w_t = (c_bbox[:, 2] - c_bbox[:, 0] + 1)
        b_h_t = (c_bbox[:, 3] - c_bbox[:, 1] + 1)
        mask = mask & (np.minimum(b_w_t, b_h_t) > 4.0) # min_feat_size>1


        c_bbox = c_bbox[mask].copy()
        c_labels = c_labels[mask].copy()
        if c_bbox.shape[0]==0:
            continue
        c_bbox[:, 0::2] = np.clip(c_bbox[:, 0::2], 0, out_size[0]-1)
        c_bbox[:, 1::2] = np.clip(c_bbox[:, 1::2], 0, out_size[1]-1)
        if landms is not None:
            c_landms = landms[mask].copy()
            for i in range(c_landms.shape[0]):
                landm = -np.ones(10)
                for k in range(5):
                    if (c_landms[i][k]>0).all():
                        pt = affine_transform(c_landms[i, k*2:(k+1)*2], trans_input)
                        if 0<pt[0]<out_size[0]-1 and 0<pt[1]<out_size[1]-1:
                            landm[k*2:(k+1)*2] = pt
                c_landms[i] = landm

        c_image = cv2.warpAffine(image, trans_input, (out_size[0], out_size[1]), flags=cv2.INTER_LINEAR)
        return c_image, c_bbox, c_labels, c_landms
    
    image, bbox, labels, landms = CenterNet_Val_Crop(image, bbox, labels, out_size, landms)
    return image, bbox, labels, landms    
    



def CenterNet_Val_Crop(image, bbox, labels, out_size, landms=None):
    h, w = image.shape[:2]
    s = max(h, w) * 1.0
    c = np.array([w / 2., h / 2.], dtype=np.float32)
    trans_input = get_affine_transform(c, s, 0, out_size)
    for i in range(bbox.shape[0]):
        bbox[i, :2] = affine_transform(bbox[i, :2], trans_input)
        bbox[i, 2:4] = affine_transform(bbox[i, 2:4], trans_input)

    centers = (bbox[:, :2] + bbox[:, 2:]) / 2
    # center in image
    mask = np.logical_and(0 < centers, centers < np.array([out_size[0], out_size[1]])).all(axis=1)
    b_w_t = (bbox[:, 2] - bbox[:, 0] + 1)
    b_h_t = (bbox[:, 3] - bbox[:, 1] + 1)
    mask = mask & (np.minimum(b_w_t, b_h_t) > 4.0)

    bbox = bbox[mask].copy()
    labels = labels[mask].copy()

    bbox[:, :2] = np.clip(bbox[:, :2], 0, out_size[0] - 1)
    bbox[:, 2:4] = np.clip(bbox[:, 2:4], 0, out_size[1] - 1)
    if landms is not None:
        landms = landms[mask].copy()
        for i in range(landms.shape[0]):
            landm = -np.ones(10)
            for k in range(5):
                if (landms[i][k]>0).all():
                    pt = affine_transform(landms[i, k*2:(k+1)*2], trans_input)
                    if 0<pt[0]<out_size[0]-1 and 0<pt[1]<out_size[1]-1:
                        landm[k*2:(k+1)*2] = pt
            landms[i] = landm     

    image = cv2.warpAffine(image, trans_input, (out_size[0], out_size[1]), flags=cv2.INTER_LINEAR)
    return image, bbox, labels, landms
    
    

def CenterNet_Transform(image, targets, out_size, split='train'):
    bbox = targets[:, :4].copy()
    landms = targets[:, 4:-1].copy() if targets.shape[1]>5 else None
    labels = targets[:, -1].copy()

    if split=='train':
        image, bbox, labels, landms = CenterNet_Train_Crop(image, bbox, labels, out_size, landms, 50)
        if random.random()<0.5:
            image, bbox, landms = HorizontalFlipAug(image, bbox, landms)
        if random.random()<0.7:
            image = RandomJitterAug(image)
    else:
        image, bbox, labels, landms = CenterNet_Val_Crop(image)
    # print(bbox.shape, landms.shape, labels[:, None].shape)
    targets = np.hstack((bbox, landms, labels[:, None]))
    return image, targets

